package com.hsxxherp.simplify.framework.algorithm;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

public class AESUtils {
    // AES 算法名称
    private static final String ALGORITHM = "AES";
    // 默认的 AES 密钥长度
    private static final int DEFAULT_KEY_SIZE = 128;

    private AESUtils() {
        // 私有构造函数，防止实例化
    }

    /**
     * 生成 AES 密钥
     *
     * @param keySize 密钥长度，可以是 128、192 或 256
     * @return 生成的密钥的 Base64 编码字符串
     * @throws Exception 生成密钥过程中可能抛出的异常
     */
    public static String generateKey(int keySize) throws Exception {
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(keySize, new SecureRandom());
        SecretKey secretKey = keyGen.generateKey();
        return Base64.getEncoder().encodeToString(secretKey.getEncoded());
    }

    /**
     * ECB 模式加密
     *
     * @param plaintext 明文数据
     * @param key       Base64 编码的密钥
     * @return 加密后的 Base64 编码字符串
     * @throws Exception 加密过程中可能抛出的异常
     */
    public static String encryptECB(String plaintext, String key) throws Exception {
        byte[] keyBytes = Base64.getDecoder().decode(key);
        SecretKeySpec secretKeySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        Cipher cipher = Cipher.getInstance(ALGORITHM + "/ECB/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        byte[] encryptedBytes = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    /**
     * ECB 模式解密
     *
     * @param ciphertext 加密后的 Base64 编码字符串
     * @param key        Base64 编码的密钥
     * @return 解密后的明文数据
     * @throws Exception 解密过程中可能抛出的异常
     */
    public static String decryptECB(String ciphertext, String key) throws Exception {
        byte[] keyBytes = Base64.getDecoder().decode(key);
        SecretKeySpec secretKeySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        Cipher cipher = Cipher.getInstance(ALGORITHM + "/ECB/PKCS5Padding");
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        byte[] decodedBytes = Base64.getDecoder().decode(ciphertext);
        byte[] decryptedBytes = cipher.doFinal(decodedBytes);
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    /**
     * CBC 模式加密
     *
     * @param plaintext 明文数据
     * @param key       Base64 编码的密钥
     * @param iv        初始化向量的 Base64 编码字符串
     * @return 加密后的 Base64 编码字符串
     * @throws Exception 加密过程中可能抛出的异常
     */
    public static String encryptCBC(String plaintext, String key, String iv) throws Exception {
        byte[] keyBytes = Base64.getDecoder().decode(key);
        byte[] ivBytes = Base64.getDecoder().decode(iv);
        SecretKeySpec secretKeySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);
        Cipher cipher = Cipher.getInstance(ALGORITHM + "/CBC/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
        byte[] encryptedBytes = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    /**
     * CBC 模式解密
     *
     * @param ciphertext 加密后的 Base64 编码字符串
     * @param key        Base64 编码的密钥
     * @param iv         初始化向量的 Base64 编码字符串
     * @return 解密后的明文数据
     * @throws Exception 解密过程中可能抛出的异常
     */
    public static String decryptCBC(String ciphertext, String key, String iv) throws Exception {
        byte[] keyBytes = Base64.getDecoder().decode(key);
        byte[] ivBytes = Base64.getDecoder().decode(iv);
        SecretKeySpec secretKeySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);
        Cipher cipher = Cipher.getInstance(ALGORITHM + "/CBC/PKCS5Padding");
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);
        byte[] decodedBytes = Base64.getDecoder().decode(ciphertext);
        byte[] decryptedBytes = cipher.doFinal(decodedBytes);
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    /**
     * 生成随机的初始化向量（IV）
     *
     * @return 初始化向量的 Base64 编码字符串
     */
    public static String generateIV() {
        byte[] iv = new byte[16];
        new SecureRandom().nextBytes(iv);
        return Base64.getEncoder().encodeToString(iv);
    }

    public static void main(String[] args) {
        try {
            // 生成密钥
            String key = generateKey(DEFAULT_KEY_SIZE);
            System.out.println("Generated Key: " + key);

            // 生成初始化向量（仅用于 CBC 模式）
            String iv = generateIV();
            System.out.println("Generated IV: " + iv);

            String plaintext = "Hello, AES!";

            // ECB 模式加密和解密
            String encryptedECB = encryptECB(plaintext, key);
            System.out.println("ECB Encrypted: " + encryptedECB);
            String decryptedECB = decryptECB(encryptedECB, key);
            System.out.println("ECB Decrypted: " + decryptedECB);

            // CBC 模式加密和解密
            String encryptedCBC = encryptCBC(plaintext, key, iv);
            System.out.println("CBC Encrypted: " + encryptedCBC);
            String decryptedCBC = decryptCBC(encryptedCBC, key, iv);
            System.out.println("CBC Decrypted: " + decryptedCBC);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}